from config import args
from math import cos, pi


def get_interval(epoch, T_0, T_mult):

    for i in range(100):
        epoch -= T_0
        T_next = T_0 * T_mult
        if epoch < T_next:
            return epoch, T_next
        elif epoch > T_next:
            T_0 = T_next
        elif epoch == T_next:
            return 0, T_next


def get_scheduler(warm_epoch=args.warm_up1, max_lr=args.lr1, T_0=args.T_01, T_mult=args.T_mult1, eta_min=args.eta_min1):
    def warm_up_cosine_annealing_warm_restart(epoch):
        if epoch <= warm_epoch:
            lr = eta_min + (max_lr - eta_min) * epoch / warm_epoch
        else:
            if T_mult == 1 or epoch < (warm_epoch + T_0):
                T = T_0
                interval = (epoch - warm_epoch) % T_0
            else:
                interval, T = get_interval(epoch - warm_epoch, T_0=T_0, T_mult=T_mult)
            lr = eta_min + 0.5 * (max_lr - eta_min) * (1 + cos(interval * pi / T))
        return lr

    return warm_up_cosine_annealing_warm_restart


if __name__ == '__main__':
    lrs = []
    scheduler = get_scheduler()
    for i in range(100):
        lrs.append(scheduler(i))
    print(lrs)
